/*
 *
 *  * Copyright (c) iwindplus Technologies Co., Ltd.2024-2030, All rights reserved.
 *
 *
 */

package com.iwindplus.base.web;

import com.iwindplus.base.domain.enums.BizCodeEnum;
import com.iwindplus.base.domain.enums.ThreadRejectedExecutionEnum;
import com.iwindplus.base.domain.exception.BizException;
import com.iwindplus.base.web.domain.property.ThreadPoolProperty;
import jakarta.annotation.Resource;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.task.TaskExecutionProperties;
import org.springframework.boot.autoconfigure.task.TaskSchedulingProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.boot.task.ThreadPoolTaskExecutorBuilder;
import org.springframework.boot.task.ThreadPoolTaskExecutorCustomizer;
import org.springframework.boot.task.ThreadPoolTaskSchedulerBuilder;
import org.springframework.boot.task.ThreadPoolTaskSchedulerCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.task.TaskDecorator;
import org.springframework.core.task.support.ContextPropagatingTaskDecorator;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

/**
 * 线程池配置.
 *
 * @author zengdegui
 * @since 2023/08/29 22:24
 */
@Slf4j
@EnableAsync
@EnableScheduling
@Configuration(proxyBeanMethods = false)
@EnableConfigurationProperties(ThreadPoolProperty.class)
@ConditionalOnProperty(prefix = "spring.task", name = "enabled", havingValue = "true", matchIfMissing = true)
public class ThreadPoolConfiguration {

    @Resource
    private TaskExecutionProperties taskExecutionProperties;

    @Resource
    private TaskSchedulingProperties taskSchedulingProperties;

    @Resource
    private ThreadPoolProperty threadPoolProperty;

    /**
     * 创建 ContextPropagatingTaskDecorator.
     *
     * @return ContextPropagatingTaskDecorator
     */
    @Bean
    public TaskDecorator contextPropagatingTaskDecorator() {
        final ContextPropagatingTaskDecorator taskDecorator = new ContextPropagatingTaskDecorator();
        log.info("ContextPropagatingTaskDecorator={}", taskDecorator);
        return taskDecorator;
    }

    /**
     * 创建 ThreadPoolTaskExecutor.
     *
     * @param builder                         builder
     * @param contextPropagatingTaskDecorator contextPropagatingTaskDecorator
     * @return ThreadPoolTaskExecutor
     */
    @Primary
    @Bean
    @ConditionalOnProperty(prefix = "spring.task", name = "enabled-execution", havingValue = "true", matchIfMissing = true)
    public ThreadPoolTaskExecutor threadPoolTaskExecutor(ThreadPoolTaskExecutorBuilder builder, TaskDecorator contextPropagatingTaskDecorator) {
        final TaskExecutionProperties.Pool pool = this.taskExecutionProperties.getPool();
        final TaskExecutionProperties.Shutdown shutdown = this.taskExecutionProperties.getShutdown();
        ThreadPoolTaskExecutor executor = builder
            .corePoolSize(Optional.ofNullable(pool.getCoreSize()).orElse(Runtime.getRuntime().availableProcessors()))
            .maxPoolSize(Optional.ofNullable(pool.getMaxSize()).orElse(Runtime.getRuntime().availableProcessors() * 2))
            .queueCapacity(Optional.ofNullable(pool.getQueueCapacity()).orElse(512))
            .keepAlive(Optional.ofNullable(pool.getKeepAlive()).orElse(Duration.ofSeconds(60)))
            .allowCoreThreadTimeOut(Optional.ofNullable(pool.isAllowCoreThreadTimeout()).orElse(Boolean.TRUE))
            .awaitTermination(Optional.ofNullable(shutdown.isAwaitTermination()).orElse(Boolean.TRUE))
            .awaitTerminationPeriod(Optional.ofNullable(shutdown.getAwaitTerminationPeriod()).orElse(Duration.ofSeconds(60)))
            .threadNamePrefix(Optional.ofNullable(this.taskExecutionProperties.getThreadNamePrefix()).orElse("task-executor-"))
            .taskDecorator(contextPropagatingTaskDecorator)
            .customizers(threadPoolTaskExecutorCustomizer())
            .build();
        log.info("ThreadPoolTaskExecutor={}", executor);
        return executor;
    }

    /**
     * 创建 ThreadPoolTaskScheduler.
     *
     * @param builder                         builder
     * @param contextPropagatingTaskDecorator contextPropagatingTaskDecorator
     * @return ThreadPoolTaskScheduler
     */
    @Primary
    @Bean
    @ConditionalOnProperty(prefix = "spring.task", name = "enabled-scheduler", havingValue = "true", matchIfMissing = true)
    public ThreadPoolTaskScheduler threadPoolTaskScheduler(ThreadPoolTaskSchedulerBuilder builder, TaskDecorator contextPropagatingTaskDecorator) {
        final TaskSchedulingProperties.Pool pool = this.taskSchedulingProperties.getPool();
        final TaskSchedulingProperties.Shutdown shutdown = this.taskSchedulingProperties.getShutdown();
        ThreadPoolTaskScheduler scheduler = builder
            .poolSize(Optional.ofNullable(pool.getSize()).orElse(Runtime.getRuntime().availableProcessors()))
            .threadNamePrefix(Optional.ofNullable(this.taskExecutionProperties.getThreadNamePrefix()).orElse("task-scheduler-"))
            .awaitTermination(Optional.ofNullable(shutdown.isAwaitTermination()).orElse(Boolean.TRUE))
            .awaitTerminationPeriod(Optional.ofNullable(shutdown.getAwaitTerminationPeriod()).orElse(Duration.ofSeconds(60)))
            .customizers(taskSchedulerCustomizer())
            .build();
        log.info("ThreadPoolTaskScheduler={}", scheduler);
        return scheduler;
    }

    /**
     * 获取拒绝策略.
     *
     * @param handlerType 拒绝策略类型
     * @return 拒绝策略
     */
    private RejectedExecutionHandler getRejectedExecutionHandler(ThreadRejectedExecutionEnum handlerType) {
        switch (handlerType) {
            case ABORT_POLICY:
                return new ThreadPoolExecutor.AbortPolicy();
            case CALLER_RUNS_POLICY:
                return new ThreadPoolExecutor.CallerRunsPolicy();
            case DISCARD_POLICY:
                return new ThreadPoolExecutor.DiscardPolicy();
            case DISCARD_OLDEST_POLICY:
                return new ThreadPoolExecutor.DiscardOldestPolicy();
            default:
                throw new BizException(BizCodeEnum.UNSUPPORTED_THREAD_REJECTED_EXECUTION, new Object[]{handlerType.getValue()});
        }
    }

    private ThreadPoolTaskExecutorCustomizer threadPoolTaskExecutorCustomizer() {
        return task -> task.setRejectedExecutionHandler(getRejectedExecutionHandler(threadPoolProperty.getRejectedExecutionHandler()));
    }

    private ThreadPoolTaskSchedulerCustomizer taskSchedulerCustomizer() {
        return scheduler -> {
            scheduler.setErrorHandler(t -> log.error("调度任务未捕获异常", t));
            scheduler.setRejectedExecutionHandler(getRejectedExecutionHandler(threadPoolProperty.getRejectedExecutionHandler()));
        };
    }
}